from numpy import (
    matmul,  # 矩阵乘法
    zeros,  # 创建全零数组
    log2,  # 计算以2为底的对数
)


def calc_sum_rate(
        channel_state,  # 信道状态信息矩阵，大小为 (用户数量, 天线数量)
        w_precoder,  # 预编码矩阵，大小为 (天线数量, 用户数量)
        noise_power_watt  # 噪声功率，单位为瓦特
) -> float:
    """
    计算系统总速率，包括用户间干扰（IUI）的影响。

    参数:
    - channel_state: 信道状态信息矩阵，表示每个用户与卫星之间的信道信息。
    - w_precoder: 预编码矩阵，表示将信号预编码以提高信号质量。
    - noise_power_watt: 噪声功率，单位为瓦特。

    返回:
    - float: 包括用户间干扰的系统总速率。
    """

    user_nr = channel_state.shape[0]  # 用户数量

    # 初始化信噪比数组
    sinr_users = zeros(user_nr)

    for user_idx in range(user_nr):
        # 提取第 user_idx 个用户的信道向量
        channel_user_H_k = channel_state[user_idx, :]

        # 计算第 user_idx 个用户的信号功率
        power_fading_precoded_sigma_x = abs(matmul(channel_user_H_k, w_precoder[:, user_idx]))**2

        # 计算其他用户对当前用户的干扰功率
        power_fading_precoded_other_users_sigma_i = [
            abs(matmul(channel_user_H_k, w_precoder[:, other_user_idx]))**2
            for other_user_idx in range(user_nr) if other_user_idx != user_idx
        ]

        # 计算总的干扰功率
        sum_power_fading_precoded_other_users_sigma_int = sum(power_fading_precoded_other_users_sigma_i)

        # 计算第 user_idx 个用户的信噪比 (SINR)
        sinr_users[user_idx] = (
            power_fading_precoded_sigma_x / (noise_power_watt + sum_power_fading_precoded_other_users_sigma_int)
        )

    # 计算每个用户的信息速率
    info_rate_users = log2(1 + sinr_users)

    # 计算所有用户信息速率的总和，即系统总速率
    sum_rate = sum(info_rate_users)

    return sum_rate
